package org.aeon.fasttasks.core.scanner;

import org.aeon.fasttasks.core.annotation.EnableTask;
import org.aeon.fasttasks.core.util.Collections;
import org.springframework.beans.BeansException;
import org.springframework.beans.factory.annotation.AnnotatedBeanDefinition;
import org.springframework.beans.factory.config.BeanDefinition;
import org.springframework.beans.factory.config.ConfigurableListableBeanFactory;
import org.springframework.beans.factory.config.ConstructorArgumentValues;
import org.springframework.beans.factory.config.ConstructorArgumentValues.ValueHolder;
import org.springframework.beans.factory.support.BeanDefinitionRegistry;
import org.springframework.beans.factory.support.BeanDefinitionRegistryPostProcessor;
import org.springframework.beans.factory.support.RootBeanDefinition;
import org.springframework.boot.autoconfigure.AutoConfigurationPackages;
import org.springframework.context.ResourceLoaderAware;
import org.springframework.context.annotation.ClassPathScanningCandidateComponentProvider;
import org.springframework.core.io.ResourceLoader;
import org.springframework.core.type.filter.AnnotationTypeFilter;

import java.lang.annotation.Annotation;
import java.util.Collection;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author aeon
 * @createTime 2025/08/11  16:11
 * @description
 */
@SuppressWarnings("unchecked")
public class AnnotationScanner implements BeanDefinitionRegistryPostProcessor, ResourceLoaderAware {

    private ResourceLoader resourceLoader;

    @Override
    public void postProcessBeanDefinitionRegistry(BeanDefinitionRegistry registry) throws BeansException {
        List<String> basePackages = AutoConfigurationPackages.get((ConfigurableListableBeanFactory) registry);
        Map<String, Class<?>> classCache = new ConcurrentHashMap<>();
        // 扫描被 EnableTask 注释的下级注释类
        ClassPathScanningCandidateComponentProvider annotationScanner =
                findAnnotationsWithMetaScanner(List.of(EnableTask.class));
        Collection<? extends Class<? extends Annotation>> annotationClasses = basePackages
                .parallelStream()
                .flatMap(basePackage -> annotationScanner.findCandidateComponents(basePackage).stream()
                        .map(beanDef -> {
                            String className = beanDef.getBeanClassName();
                            Class<?> clazz = classCache.computeIfAbsent(className, name -> {
                                try {
                                    return Class.forName(name, false,
                                            Thread.currentThread().getContextClassLoader());
                                } catch (ClassNotFoundException e) {
                                    throw new RuntimeException(e);
                                }
                            });
                            return clazz.isAnnotation() ? (Class<? extends Annotation>) clazz : null;
                        })
                        .filter(Objects::nonNull)
                )
                .distinct()
                .toList();
        // 扫描被 EnableTask 注释的类
        ClassPathScanningCandidateComponentProvider classesScanner = findClassesScanner(
                Collections.of(annotationClasses, EnableTask.class));
        Collection<? extends Class<?>> classes = basePackages.parallelStream()
                .flatMap(basePackage -> classesScanner.findCandidateComponents(basePackage)
                        .stream()
                        .map(BeanDefinition::getBeanClassName)
                        .map(className -> classCache.computeIfAbsent(className, name -> {
                            try {
                                return Class.forName(name, false,
                                        Thread.currentThread().getContextClassLoader());
                            } catch (ClassNotFoundException e) {
                                throw new RuntimeException(e);
                            }
                        }))
                )
                .toList();
        // 注册扫描结果
        if (registry.containsBeanDefinition(TaskScanResult.NAME)) {
            BeanDefinition bd = registry.getBeanDefinition(TaskScanResult.NAME);
            ConstructorArgumentValues args = bd.getConstructorArgumentValues();
            ValueHolder valueHolder = args.getGenericArgumentValue(Collection.class);
            if (valueHolder != null && valueHolder.getValue() instanceof Collection) {
                ((Collection<Class<?>>) valueHolder.getValue()).addAll(classes);
            }
        } else {
            RootBeanDefinition def = new RootBeanDefinition(TaskScanResult.class);
            def.getConstructorArgumentValues().addGenericArgumentValue(classes);
            registry.registerBeanDefinition(TaskScanResult.NAME, def);
        }
    }

    /**
     * 扫描指定包，返回被指定元注解标注的注解类，
     *
     * @param metaAnnotationClass 元注解类型
     * @return Collection<? extends Class < ? extends Annotation>>
     */
    private ClassPathScanningCandidateComponentProvider findAnnotationsWithMetaScanner(
            List<Class<? extends Annotation>> metaAnnotationClass) {
        ClassPathScanningCandidateComponentProvider provider =
                new ClassPathScanningCandidateComponentProvider(false) {
                    @Override
                    protected boolean isCandidateComponent(AnnotatedBeanDefinition beanDefinition) {
                        return true;
                    }
                };
        provider.setResourceLoader(resourceLoader);
        for (Class<? extends Annotation> annotation : metaAnnotationClass) {
            provider.addIncludeFilter(new AnnotationTypeFilter(annotation));
        }
        return provider;
    }

    /**
     * 扫描指定包，返回被指定注解标注的类
     *
     * @param annotationClass 注解类型
     * @return
     */
    private ClassPathScanningCandidateComponentProvider findClassesScanner(
            Collection<? extends Class<? extends Annotation>> annotationClass) {
        ClassPathScanningCandidateComponentProvider provider =
                new ClassPathScanningCandidateComponentProvider(false);
        provider.setResourceLoader(resourceLoader);
        for (Class<? extends Annotation> annotation : annotationClass) {
            provider.addIncludeFilter(new AnnotationTypeFilter(annotation));
        }
        return provider;
    }


    @Override
    public void setResourceLoader(ResourceLoader resourceLoader) {
        this.resourceLoader = resourceLoader;
    }
}
